﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Runtime.InteropServices;
using SharpML.Api.Enumerations;
using OpenCL.Net;


namespace SharpML.Api.Extensions
{
    /// <summary>
    /// Originally borrowed from OpenClooVision - http://opencloovision.codeplex.com/ in SharpRBM
    /// Then borrowed from that from SharpRBM was rewritten and refactored into SharpML
    /// </summary>
    public static class OpenCLExtensions
    {
        private static int _intPtrSize = -1;
        public static int IntPtrSize
        {
            get
            {
                if (_intPtrSize == -1)
                {
                    Marshal.SizeOf( typeof( IntPtr ) );
                }
                return _intPtrSize;
            }
        }

        public static Kernel CreateKernel(this Program program, string kernelName)
        {
            ErrorCode error;
            Kernel kernel = Cl.CreateKernel(program, kernelName, out error);
            if (error != ErrorCode.Success)
            {
                throw new InvalidProgramException(string.Format("Unable to run CreateKernel for {1}: {0}!", error, kernelName));
            }
            return kernel;
        }

        public static Kernel SetKernelArg<T>( this Kernel kernel, uint argId, T value )
        {
            GCHandle gcHandle = GCHandle.Alloc( value, GCHandleType.Pinned );
            try
            {
                ErrorCode error = Cl.SetKernelArg( kernel, argId, new IntPtr( Marshal.SizeOf( typeof( T ) ) ), gcHandle.AddrOfPinnedObject() );

                if( error != ErrorCode.Success )
                {
                    throw new InvalidProgramException( string.Format( "Unable to run SetKernelArg<T>: {0}!", error ) );
                }
            }
            finally
            {
                gcHandle.Free();
            }

            return kernel;
        }

        public static Kernel SetKernelArg(this Kernel kernel, uint argId, IMem mem, int size)
        {
            ErrorCode error = Cl.SetKernelArg(kernel, argId, new IntPtr(size), mem);
            if (error != ErrorCode.Success)
            {
                throw new InvalidProgramException(string.Format("Unable to run SetKernelArg: {0}!", error));
            }
            return kernel;
        }

        public static Kernel SetKernelArgValue( this Kernel kernel, uint argId, IMem mem )
        {
            return kernel.SetKernelArg( argId, mem, sizeof( float ) );
        }

        public static Kernel SetKernelArgFloat(this Kernel kernel, uint argId, IMem mem)
        {
            return kernel.SetKernelArg(argId, mem, sizeof(float));
        }

        public static Kernel SetKernelArgBool(this Kernel kernel, uint argId, IMem mem)
        {
            return kernel.SetKernelArg(argId, mem, sizeof(Bool));
        }

        public static void EnqueueNDRangeKernel(this Kernel kernel, CommandQueue commandQueue, params int[] sizes)
        {
            Event clEvent;
            ErrorCode error = Cl.EnqueueNDRangeKernel(
                commandQueue,
                kernel,
                (uint)sizes.Length,
                null,
                sizes.ToList().Select(size => new IntPtr(size)).ToArray(),
                null,
                0,
                null,
                out clEvent);
            if (error != ErrorCode.Success)
            {
                throw new InvalidOperationException("EnqueueNDRangeKernel failed:" + error);
            }
            clEvent.Dispose();
        }

        public static void EnqueueTask(this Kernel kernel, CommandQueue commandQueue)
        {
            Event clEvent;
            ErrorCode error = Cl.EnqueueTask(
                commandQueue,
                kernel,
                0,
                null,
                out clEvent);
            if (error != ErrorCode.Success)
            {
                throw new InvalidOperationException("EnqueueTask failed:" + error);
            }
            clEvent.Dispose();
        }

        public static void Read(this IMem mem, CommandQueue commandQueue, int size, object data, BlockModes blockMode)
        {
            Event clEvent;
            IntPtr event_handle = IntPtr.Zero;
            ErrorCode error = Cl.EnqueueReadBuffer(commandQueue, mem, blockMode == BlockModes.Blocking ? Bool.True : Bool.False, IntPtr.Zero, new IntPtr(size), data, 0, null, out clEvent);
            if (error != ErrorCode.Success)
            {
                throw new InvalidOperationException("EnqueueReadBuffer failed for _weightsIMem:" + error);
            }
            clEvent.Dispose();
        }

        public static void ReadFloatArray(this IMem mem, CommandQueue commandQueue, float[] array, BlockModes blockMode)
        {
            mem.Read(commandQueue, sizeof(float) * array.Length, array, blockMode);
        }

        public static void Write(this IMem mem, CommandQueue commandQueue, int size, object data, BlockModes blockMode)
        {
            Event clEvent;
            ErrorCode error = Cl.EnqueueWriteBuffer(commandQueue, mem, blockMode == BlockModes.Blocking ? Bool.True : Bool.False, IntPtr.Zero, new IntPtr(size), data, 0, null, out clEvent);
            if (error != ErrorCode.Success)
            {
                throw new InvalidProgramException(string.Format("EnqueueWriteBuffer failed: {0}!", error));
            }
            clEvent.Dispose();
        }

        public static void WriteFloatArray(this IMem mem, CommandQueue commandQueue, float[] array, BlockModes blockMode)
        {
            mem.Write(commandQueue, sizeof(float) * array.Length, array, blockMode);
        }


        public static IMem AllocateFloatArray(this Context context, float[] array, MemFlags flags)
        {
            ErrorCode error;
            int size = sizeof(float) * array.Count();
            IMem result = Cl.CreateBuffer(context, flags, (IntPtr)size, array, out error);
            if (error != ErrorCode.Success)
            {
                throw new InvalidProgramException(string.Format("CreateBuffer failed:", error));
            }
            return result;
        }

        public static IMem AllocateIntArray(this Context context, int[] array, MemFlags flags)
        {
            ErrorCode error;
            int size = 4 * array.Count();
            IMem result = Cl.CreateBuffer(context, flags, (IntPtr)size, array, out error);
            if (error != ErrorCode.Success)
            {
                throw new InvalidProgramException(string.Format("CreateBuffer failed:", error));
            }
            return result;
        }

        public static void Finish(this CommandQueue commandQueue)
        {
            ErrorCode error = Cl.Finish(commandQueue);
            if (error != ErrorCode.Success)
            {
                throw new InvalidProgramException(string.Format("Finish failed:", error));
            }
        }
        /*public static void SetIMemoryArgument(this Kernel kernel, int argumentId, object value)
        {
            ErrorCode error = SetKernelArg(kernel, 0, new IntPtr(inPtrSize), hDeviceIMemA);
            Assert.AreEqual(ErrorCode.Success, error);
        }

        public static void SetArguments(this Kernel kernel, params object[] args)
        {            
            for (int i = 0; i < args.Length; i++)
            {
                try
                {
                    object arg = args[i];
                    if (arg == null) { kernel.SetIMemoryArgument(i, null); continue; }
                    var typeCode = Type.GetTypeCode(arg.GetType());
                    switch (typeCode)
                    {
                        //case TypeCode.Boolean:
                        //    break;
                        case TypeCode.Byte:
                            kernel.SetValueArgument<byte>(i, (byte)arg);
                            break;
                        case TypeCode.Char:
                            kernel.SetValueArgument<char>(i, (char)arg);
                            break;
                        //case TypeCode.DBNull:
                        //    break;
                        //case TypeCode.DateTime:
                        //    break;
                        case TypeCode.Decimal:
                            kernel.SetValueArgument<decimal>(i, (decimal)arg);
                            break;
                        case TypeCode.Double:
                            kernel.SetValueArgument<double>(i, (double)arg);
                            break;
                        case TypeCode.Empty:
                            kernel.SetIMemoryArgument(i, null);
                            break;
                        case TypeCode.Int16:
                            kernel.SetValueArgument<short>(i, (short)arg);
                            break;
                        case TypeCode.Int32:
                            kernel.SetValueArgument<int>(i, (int)arg);
                            break;
                        case TypeCode.Int64:
                            kernel.SetValueArgument<long>(i, (long)arg);
                            break;
                        case TypeCode.Object:
                            {
                                ClooBuffer mem = arg as ClooBuffer;
                                if (mem != null) { kernel.SetIMemoryArgument(i, mem.DeviceBuffer); continue; }
                            }
                            {
                                ComputeIMemory mem = arg as ComputeIMemory;
                                if (mem != null) { kernel.SetIMemoryArgument(i, mem); continue; }
                            }
                            goto default;
                        case TypeCode.SByte:
                            kernel.SetValueArgument<sbyte>(i, (sbyte)arg);
                            break;
                        case TypeCode.Single:
                            kernel.SetValueArgument<float>(i, (float)arg);
                            break;
                        //case TypeCode.String:
                        //    break;
                        case TypeCode.UInt16:
                            kernel.SetValueArgument<ushort>(i, (ushort)arg);
                            break;
                        case TypeCode.UInt32:
                            kernel.SetValueArgument<uint>(i, (uint)arg);
                            break;
                        case TypeCode.UInt64:
                            kernel.SetValueArgument<ulong>(i, (ulong)arg);
                            break;
                        default:
                            throw new NotSupportedException("Type " + arg.GetType().ToString() + " is not supported as argument");
                    }
                }
                catch (InvalidKernelArgumentsComputeException)
                {
                    throw new ArgumentException(String.Format("Invalid kernel parameter specified on index {0} for function {1}", i, kernel.FunctionName));
                }
            }
        }*/
    }
}
